package ace.cmp.util;

import ace.cmp.util.model.dto.TreeNode;
import java.util.AbstractMap.SimpleEntry;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;
import org.apache.commons.collections.CollectionUtils;

/**
 * @author caspar
 * @date 2024/7/18 11:21
 */
public final class TreeUtils {

  public static <T, TId> List<TreeNode<T>> buildTree(
      Collection<T> values,
      Function<T, Boolean> isRootFn,
      Function<T, TId> getIdFn,
      Function<T, List<T>> getChildrenFn) {

    Map<TId, TreeNode<T>> nodeMap = new HashMap<>();
    for (T value : values) {
      TreeNode<T> node = TreeNode.<T>builder().value(value).build();
      TId id = getIdFn.apply(value);
      nodeMap.put(id, node);
    }

    for (T value : values) {
      TId id = getIdFn.apply(value);
      TreeNode<T> node = nodeMap.get(id);
      List<T> children = getChildrenFn.apply(value);
      if (CollectionUtils.isEmpty(children)) {
        continue;
      }
      List<TreeNode<T>> nodeChildren =
          children.stream()
              .map(p -> nodeMap.get(getIdFn.apply(p)))
              .filter(Objects::nonNull)
              .collect(Collectors.toList());
      node.addChildren(nodeChildren);
    }

    return nodeMap.values().stream().filter(p -> isRootFn.apply(p.getValue())).collect(Collectors.toList());
  }

  public static <T> void traversalSort(
      List<TreeNode<T>> treeNodes, Comparator<TreeNode<T>> comparator) {
    treeNodes.sort(comparator);
    TreeUtils.traversal(treeNodes, treeNode -> treeNode.getChildren().sort(comparator));
  }

  /**
   * 遍历
   * @param treeNodes
   * @param consumer
   * @param <T>
   * @param <TId>
   */
  public static <T, TId> void traversal(
      List<TreeNode<T>> treeNodes, Consumer<TreeNode<T>> consumer) {
    for (TreeNode<T> treeNode : treeNodes) {
      consumer.accept(treeNode);
      traversal(treeNode.getChildren(), consumer);
    }
  }

  /**
   * 从下到上遍历,并过滤重复id.
   *
   * @param bottomPredicate 判断是否最底层节点
   * @param getIdFn 获取id函数
   * @param getParentIdFn 获取父id函数
   * @param consumer 遍历消费者 key=level 从下至上的层数 value=节点数据
   */
  public static <T, TId> void traversalBottomToUp(
      List<T> values,
      Predicate<T> bottomPredicate,
      Function<T, TId> getIdFn,
      Function<T, TId> getParentIdFn,
      Consumer<Entry<Long, T>> consumer
  ) {

    Map<TId, T> currentValues = values
        .stream()
        .filter(bottomPredicate)
        .collect(
            HashMap::new,
            (m, entry) -> {
              TId id = getIdFn.apply(entry);
              m.put(id, entry);
            },
            HashMap::putAll);
    Map<TId, T> tmpParentValues = new HashMap<>();
    long level = 0;
    while (CollectionUtils.isNotEmpty(currentValues.entrySet())) {
      for (Entry<TId, T> entry : currentValues.entrySet()) {
        T value = entry.getValue();
        SimpleEntry<Long, T> eventParams = new SimpleEntry<>(level, value);
        consumer.accept(eventParams);

        T parentValue = values.stream()
            .filter(parent -> getParentIdFn.apply(value).equals(getIdFn.apply(parent)))
            .findFirst()
            .orElse(null);

        if (Objects.nonNull(parentValue)) {
          TId parentId = getIdFn.apply(parentValue);
          tmpParentValues.put(parentId, parentValue);
        }
      }
      currentValues.clear();
      currentValues.putAll(tmpParentValues);
      tmpParentValues.clear();
      level++;
    }
  }
}
